package com.gt.stream

import org.apache.kafka.clients.consumer.{ConsumerConfig, ConsumerRecord}
import org.apache.spark.SparkConf
import org.apache.spark.streaming.dstream.{DStream, InputDStream}
import org.apache.spark.streaming.kafka010.{ConsumerStrategies, KafkaUtils, LocationStrategies, LocationStrategy}
import org.apache.spark.streaming.{Duration, Seconds, StreamingContext}

/**
 * SparkStreaming 读取kafka数据
 * 使用 updateStateByKey 记录状态
 * 异常退出后，从 checkpoint 恢复状态
 * 注: 从checkpoint恢复，代码是不能改动的
 */
object Streaming_WC_kafka_state_cp_03 {

  def main(args: Array[String]): Unit = {
    val ssc: StreamingContext = StreamingContext.getOrCreate(checkpoint, execute)

    ssc.start()
    ssc.awaitTermination()

  }

  val checkpoint = "data/Streaming_WC_kafka_03"

  def execute(): StreamingContext = {
    val conf: SparkConf = new SparkConf().setAppName("xx").setMaster("local[1]")
    val ssc: StreamingContext = new StreamingContext(conf, Duration(1000L))
    ssc.checkpoint(checkpoint)

    val kafkaPara: Map[String, Object] = Map[String, Object](
      ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG -> "localhost:9092",
      ConsumerConfig.GROUP_ID_CONFIG -> "gid_4_test",
      ConsumerConfig.AUTO_OFFSET_RESET_CONFIG -> "latest", // earliest
      "key.deserializer" -> "org.apache.kafka.common.serialization.StringDeserializer",
      "value.deserializer" -> "org.apache.kafka.common.serialization.StringDeserializer"
    )

    val stream: InputDStream[ConsumerRecord[String, String]] = KafkaUtils.createDirectStream[String, String](
      ssc,
      LocationStrategies.PreferConsistent,
      ConsumerStrategies.Subscribe[String, String](Set("test"), kafkaPara)
    )

    val word2one: DStream[(String, Int)] = stream.map(_.value()).flatMap(_.split(" ")).map((_, 1))

    val result: DStream[(String, Int)] = word2one.updateStateByKey(
      (seq: Seq[Int], buf: Option[Int]) => {
        Some(buf.getOrElse(0) + seq.sum)
      }
    )

    result.print()
    ssc
  }

}
